import cv2
import numpy as np
import re
import os


def get_all_files(data_path, cate, n_lim, IMG_W, IMG_H):
    """
    Get all data from files

    :param data_path: The path to the directory of files
    :return: tuple(x, y): The data x in shape (batch, height, width, channels)
    and the labels y in shape (batch, class_label)
    """
    x, y = [], []
    all_filenames = os.listdir(data_path)  # get all file names of data
    xregexp = re.compile(r'^([a-zA-Z]+)\.')  # the regular expression to extract class information of files
    n = 0
    for xfilename in all_filenames:
        # get label
        xmatch = xregexp.match(xfilename)
        if xmatch is None:
            continue
        cls = xmatch[1]
        if cls != cate:
            continue

        n += 1
        if n > n_lim:
            break

        # get data
        xdata = cv2.imread(data_path + '/' + xfilename)
        xdata = cv2.resize(xdata, [IMG_W, IMG_H])

        x.append(xdata)

    # tidy data
    x = np.array(x, dtype=np.float32) / 255.
    return x


if '__main__' == __name__:
    import matplotlib.pyplot as plt

    # ⑦	调用以上get_all_files函数，加载进来相关数据。
    data_path = r'../../../../large_data/DL1/catdog_data/data/train'
    IMG_H = 300
    IMG_W = 400

    plt.figure(figsize=[16, 8])
    spr = 4  # subplot row
    spc = 8  # subplot column
    spn = 0

    n_pics = spr * spc
    n_dog = int(np.ceil(n_pics / 2))
    n_cat = n_pics - n_dog


    def show_data(cate, n):
        global spn
        x = get_all_files(data_path, cate, n)
        for i, pic in enumerate(x):
            spn += 1
            if spn > spr * spc:
                break
            plt.subplot(spr, spc, spn)
            plt.axis('off')
            plt.title(cate)
            plt.imshow(pic)


    show_data('dog', n_dog)
    show_data('cat', n_cat)
